import sys
import os
# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

import matplotlib.pyplot as plt
import numpy as np

# 导入自定义模块
import pendulum_config
from generate_simulated_data import generate_simulated_data
from fit_model import fit_model
from plot_phase_diagram import plot_phase_diagram
from create_animation import create_animation

# 从pendulum_config模块获取参数
g = pendulum_config.g
l = pendulum_config.l
m = pendulum_config.m

def main():
    print("有空气阻力的单摆运动仿真")
    print("=" * 40)
    
    # 真实参数
    k_true = 0.2  # 真实阻尼系数 (kg/s)
    theta0 = 0.3  # 初始角度 (弧度)
    omega0 = 0.0  # 初始角速度 (弧度/秒)
    
    print(f"真实阻尼系数: k = {k_true} kg/s")
    print(f"初始条件: θ₀ = {theta0:.2f} rad, ω₀ = {omega0:.2f} rad/s")
    
    # 1. 生成模拟实验数据
    print("\n1. 生成模拟实验数据...")
    t, theta_noisy, theta_true = generate_simulated_data(k_true, theta0, omega0, t_max=20, noise_level=0.05)
    
    # 绘制实验数据
    plt.figure(figsize=(12, 6))
    plt.plot(t, theta_noisy, 'o', markersize=4, alpha=0.7, label='实验数据')
    plt.plot(t, theta_true, '-', linewidth=2, label='真实值')
    plt.xlabel('时间 t (秒)', fontsize=14)
    plt.ylabel('角度 θ (弧度)', fontsize=14)
    plt.title('模拟实验数据', fontsize=16)
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('simulated_data.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    # 2. 参数拟合
    print("\n2. 参数拟合...")
    k_fitted, k_error = fit_model(t, theta_noisy, theta0, omega0)
    print(f"拟合阻尼系数: k = {k_fitted:.4f} ± {k_error:.4f} kg/s")
    print(f"相对误差: {abs(k_fitted - k_true) / k_true * 100:.2f}%")
    
    # 3. 绘制相图
    print("\n3. 绘制相图...")
    plot_phase_diagram(k_fitted, theta0, omega0, t_max=20)
    
    # 4. 创建动画
    print("\n4. 创建动画...")
    print("正在生成动画，请稍候...")
    anim = create_animation(k_fitted, theta0, omega0, t_max=10)
    print("动画已保存为 '有阻尼单摆动画.gif'")
    
    print("\n完成！")

if __name__ == "__main__":
    main()